import os
import argparse
from pathlib import Path

# --- 要下载的模型列表 (可根据需要自行修改) ---
MODELS_TO_DOWNLOAD = [
    "BAAI/bge-large-zh-v1.5",
    "BAAI/bge-small-zh-v1.5",
    "BAAI/bge-reranker-base",
]

def setup_models(proxy: str | None = None):
    """
    一步到位，下载并安装所有必需的 Hugging Face 模型。

    本脚本会自动将模型下载并缓存到 Hugging Face 的默认全局目录，
    无需手动打包、解压或移动。
    """
    print("=" * 70)
    print("      Hugging Face 模型一键安装脚本")
    print("=" * 70)

    # 1. 设置代理 (如果提供)
    if proxy:
        print(f"[*] 步骤 1/3: 检测到代理参数，正在设置代理: {proxy}")
        os.environ['HTTP_PROXY'] = proxy
        os.environ['HTTPS_PROXY'] = proxy
    else:
        print("[*] 步骤 1/3: 未指定代理，将尝试直接连接。")

    # 2. 导入核心库
    try:
        from huggingface_hub import snapshot_download
        from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
        print(f"[*] 步骤 2/3: 核心库加载成功。")
        print(f"   - 模型将被自动安装到: {HUGGINGFACE_HUB_CACHE}")
    except ImportError:
        print("\n[错误] 核心库 'huggingface_hub' 未安装。")
        print("请先运行: pip install huggingface_hub")
        return

    # 3. 循环下载并缓存所有模型
    print("\n[*] 步骤 3/3: 开始下载所有模型...")
    print("   - Hugging Face 库将自动显示每个文件的下载进度。")
    print("-" * 70)

    for model_name in MODELS_TO_DOWNLOAD:
        print(f"\n>>> 开始处理模型: {model_name}\n")
        try:
            # snapshot_download 会自动处理缓存，如果文件已存在则会跳过。
            # 它会为大文件显示详细的下载进度条。
            snapshot_download(
                repo_id=model_name,
                local_files_only=False, # 确保联网下载
                resume_download=True,   # 支持断点续传
            )
            print(f"\n[✓] 模型 '{model_name}' 已成功安装或验证。")
        except Exception as e:
            print(f"\n[错误] 处理模型 {model_name} 时发生严重错误: {e}")
            print("请检查您的网络连接或代理设置后重试。")
            # 即使一个模型失败，也尝试继续下载下一个
            continue
    
    print("\n" + "=" * 70)
    print("🎉 恭喜！所有模型均已安装到正确位置！")
    print("=" * 70)
    print("\n现在，您可以在应用中直接使用这些模型了。")
    print("请确保 .env 文件中的 'TRANSFORMERS_OFFLINE=1' 已被启用，")
    print("以强制应用使用这些已下载的本地模型。")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="一键下载并安装 Hugging Face 模型到本地缓存。",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        "--proxy",
        type=str,
        default=None,
        help="[可选] 用于下载的 HTTP/HTTPS 代理, 例如: http://127.0.0.1:7890"
    )
    args = parser.parse_args()
    setup_models(args.proxy)
